package cn.edu.hit.kg;

import cn.edu.hit.kg.data.CypherRunner;
import cn.edu.hit.kg.data.GraphDataProvider;
import cn.edu.hit.kg.data.Node;
import org.neo4j.driver.v1.Record;
import org.neo4j.driver.v1.types.Relationship;
import org.neo4j.driver.v1.types.Type;
import org.neo4j.driver.v1.types.TypeSystem;

import java.util.*;
import java.util.concurrent.ConcurrentHashMap;
import java.util.stream.Collectors;
import java.util.stream.StreamSupport;

/**
 * Created by Wang on 2018/11/19.
 */
public class Neo4jDataProvider implements GraphDataProvider, CypherRunner {
    private Map<Long, Node> nodes = new ConcurrentHashMap<>();
    private Map<Long, Set<Long>> edges = new ConcurrentHashMap<>();
    private Map<String, Set<Long>> labelGroups = new ConcurrentHashMap<>();
    private boolean inited = false;

    static String url = "bolt://localhost";
    static String username = "neo4j";
    static String password = "123456";

    public void addNode(NodeImpl node) {
        if (!inited) {
            init(url,username,password);
        }
        this.nodes.put(node.id(), node);
        for (String label : node.labels()) {
            Set<Long> set = this.labelGroups.getOrDefault(label, new HashSet<>());
            set.add(node.id());
            this.labelGroups.put(label, set);
        }
    }

    public Neo4jDataProvider() {
    }

    public void init(String url,String username,String password) {
        if (!inited) {
            inited = true;
        } else {
            return;
        }
        try (NeoDriver driver = new NeoDriver(url,username,password)) {
            Map<Long, Set<Long>> dbEdges = new ConcurrentHashMap<>();
            Map<Long,NodeImpl> dbNodes = new ConcurrentHashMap<>();

            List<Long> faultIds = driver.query("match (n:Fault) return id(n) as nid ").list(r -> r.get("nid").asLong());
            faultIds.forEach(rootId -> {
                Map<Long, List<Long>> map = driver.query("start n=node($id) match p=(n)-[*]->(m) " +
                        "with collect(relationships(p)) as rs " +
                        "with reduce(rlist=[],rl in rs|rlist +rl) as rls " +
                        "unwind rls as rr " +
                        "return distinct id(STARTNODE(rr)) as rb,id(endNode(rr)) as rd", Collections.singletonMap("id", rootId))
                        .list().stream()
                        .collect(Collectors.groupingBy(record -> record.get("rb").asLong(),
                                Collectors.mapping(record1 -> record1.get("rd").asLong(), Collectors.toList())));
                map.forEach((fromId, toIds) -> {
                    Set<Long> toSet = dbEdges.getOrDefault(fromId, new HashSet<>());
                    toSet.addAll(toIds);
                    dbEdges.put(fromId, toSet);
                });
                Set<Long> subNodeIds = map.values().stream().flatMap(Collection::stream).collect(Collectors.toSet());
                subNodeIds.add(rootId);

                driver.query("match (n) where id(n) in $ids return n", Collections.singletonMap("ids", subNodeIds))
                        .list(record -> record.get("n").asNode()).stream().forEach(node1 -> {
                    if (!nodes.containsKey(node1.id())) {
                        NodeImpl node = new NodeImpl();
                        if(node1.asMap().containsKey("user_id")){
                            node.setId(Long.parseLong(node1.get("user_id").asString()));
                        }else {
                            node.setId(node1.id());
                        }
                        node.setLabels(StreamSupport.stream(node1.labels().spliterator(), false).collect(Collectors.toSet()));
                        node.setValueMap(node1.asMap());
                        node.getValueMap().put("db_id",node1.id());
                        this.addNode(node);
                        dbNodes.put(node1.id(),node);
                    }
                });

                dbEdges.forEach((sDbId,tDbList)->{
                    edges.put(dbNodes.get(sDbId).id(),tDbList.stream().map(tId->dbNodes.get(tId).id()).collect(Collectors.toSet()));
                });
            });
        }
    }

    @Override
    public Map<Long, Node> getNodes() {
        if (!inited) {
            init(url,username,password);
        }
        return nodes;
    }

    @Override
    public Map<Long, Set<Long>> getEdges() {
        if (!inited) {
            init(url,username,password);
        }
        return edges;
    }

    @Override
    public Map<String, Set<Long>> getLabelGroups() {
        if (!inited) {
            init(url,username,password);
        }
        return labelGroups;
    }

    @Override
    public boolean config(Map configure) {
        url = configure.getOrDefault("NEO4J_BOLT", "bolt://localhost").toString();
        username = configure.getOrDefault("USERNAME", "neo4j").toString();
        password = configure.getOrDefault("PASSWORD", "123456").toString();
        return Boolean.parseBoolean(configure.getOrDefault("useNeo4j", false).toString());
    }

    @Override
    public List<Map<String, Object>> runCyper(String cyper) {
        try (NeoDriver driver = new NeoDriver(url,username,password)) {
            return driver.execute(session -> {
                TypeSystem typeSystem = session.typeSystem();
                return session.run(cyper).list(r -> r.asMap(value -> {
                    if (typeSystem.NODE().isTypeOf(value)) {
                        org.neo4j.driver.v1.types.Node node = value.asNode();
                        Map<String, Object> nodeMap = new HashMap<>();
                        nodeMap.put("identity", node.id());
                        nodeMap.put("labels", node.labels());
                        nodeMap.put("properties", node.asMap());
                        return nodeMap;
                    }
                    if (typeSystem.RELATIONSHIP().isTypeOf(value)) {
                        Relationship relationship = value.asRelationship();
                        Map<String, Object> rmap = new HashMap<>();
                        rmap.put("start", relationship.startNodeId());
                        rmap.put("end", relationship.endNodeId());
                        rmap.put("type", relationship.type());
                        rmap.put("properties", relationship.asMap());
                        return rmap;
                    }
                    if (typeSystem.MAP().isTypeOf(value)) {
                        return value.asMap();
                    }
                    if (typeSystem.LIST().isTypeOf(value)) {
                        return value.asList();
                    }
                    if (typeSystem.NUMBER().isTypeOf(value)) {
                        return value.asNumber();
                    }
                    if (typeSystem.STRING().isTypeOf(value)) {
                        return value.asString();
                    }
                    return value.asObject();
                }));
            });
        }
    }
}
